/** \file    coherentMotion.cpp
  * \brief   Finds coherent motion of all segments.
  * \date    2001 - 2003
  * \author  Fedor Chelnokov
  * \version $Header: /home/cscom1r4/cvsroot/cs/segm/src/CoherentMotion.cpp,v 1.2 2007/08/26 13:41:42 cscom1r4 Exp $
  */

#include "segm/Segmentation.h"
#include "cinline.h"
#include <iostream>
#include <algorithm>
#include "std_min_max.h"
#include "vec2d/vec2d.h"

#include "imglib/type.h"
#include "typelib/affine.h"

using std::clog;
using std::endl;
using YATL::Arr2d;
using vec2d::Vec2i;

///adds with some offset and scale one array to another, arrays may have different sizes
static void scaledAdd( Arr2d &tgt, double scl, const Arr2d &src, int_point d, ///< offset of src relative to tgt
	double restAddend ) ///<increment for those elements of tgt that have not found addend in src
{
	Vec2i tp; //tgt place
	Vec2i sp; //src place
	for( tp.y = 0; tp.y != tgt.Height(); ++tp.y ) {
		sp.y = tp.y - d.y;
		for( tp.x = 0; tp.x != tgt.Width(); ++tp.x ) {
			sp.x = tp.x - d.x;
			if( sp.y < 0 || sp.x < 0 || sp.y >= src.Height() || sp.x >= src.Width() ) {
				tgt(tp) += scl*restAddend;
			} else {
				tgt(tp) += scl*src(sp);
			}
		}
	}
}

///sets initial approximation of displacements for all segments
static void intitialApproximation(Segmentation & s, float tau) 
{
	Arr2d maxFun;
	for( int i = 0; i != s.count(); ++i ) {
		Segment &sg = s.segmentByIndex(i);
		maxFun = sg.profile;
		IntIntMap::iterator n; //by neighbours
		for( n = sg.neighbours.begin(); n != sg.neighbours.end(); ++n ) {
			const Segment &nsg = s.segmentByIndex(n->second);
			const Arr2d &nprofile = nsg.profile;
			scaledAdd( maxFun, tau, nprofile,
				(nsg.pdisplacement()-nsg.profileCenter())-(sg.pdisplacement()-sg.profileCenter()),
				*std::min_element(nprofile.begin(),nprofile.end()) );
		}
		Point2i am = maxFun.ArgMax();
		sg.setPpos(int_point(am.x, am.y));
	}
}

///makes a step in the iteration process
static void makeIterStep( Segmentation &s, float mu ) 
{
	int i;
	//segments' displacements from previous iteration
	std::vector<int_point> old( s.count() );
	for( i = 0; i != s.count(); ++i ) {
		old[i] = s.segmentByIndex(i).adisplacement();
	}

	for( i = 0; i != s.count(); ++i ) {
		Segment &sg = s.segmentByIndex(i);

		//q(x)=ax^2+bx+c
		//q(x+1)-q(x)=a(2x+1)+b
		//a=neighbours b=-sum(2old[n->second].x);
		int a = sg.neighbours.size();
		int bx = 0;
		int by = 0;
		int sqX0 = 0;
		int sqY = 0;
		IntIntMap::iterator n; //by neighbours
		for( n = sg.neighbours.begin(); n != sg.neighbours.end(); ++n ) {
			const int_point oldDsp = sg.a2ppos(old[n->second]);
			bx -= oldDsp.x;
			by -= oldDsp.y;
			sqX0 -= sqr(oldDsp.x);
			sqY -= sqr(oldDsp.y);
		}
		bx *= 2;
		by *= 2;

		int_point mxDsp;
		double max = -DBL_MAX;
		double recip = a==0 ? 0.0 : mu / double(a);
		for( int y = 0; y != sg.profile.Height(); sqY-=a*(2*y+1)+by, ++y ) {
			int sqX = sqX0 + sqY;
			const double *p = sg.profile.BeginRow(y);
			for( int x = 0; x != sg.profile.Width(); sqX-=a*(2*x+1)+bx, ++x, ++p ) {
				double val = recip * sqX + *p;
				if( val > max ) {
					max = val;
					mxDsp.x = x;
					mxDsp.y = y;
				}
			}
		}
		assert( max != -DBL_MAX );
		sg.setPpos( mxDsp );
	}
}

///calculates current value of the functional
///\param mu in/out 
///\param beforeIter is set if its preliminary functional computation (affects mu output)
static double targetFunctional(Segmentation &s, float &mu, bool beforeIter, const MotionParams & param)
{
	double Fc = 0; //correlation part
	double Fd = 0; //displacement part
	for( int i = 0; i != s.count(); ++i ) {
		Segment &sg = s.segmentByIndex(i);
		Fc += sg.profileVal();
		IntIntMap::iterator n; //by neighbours
		double rec = 1.0 / double(sg.neighbours.size());
		for( n = sg.neighbours.begin(); n != sg.neighbours.end(); ++n ) {
			Segment &sg2 = s.segmentByIndex(n->second);
			Fd -= square(sg.adisplacement()-sg2.adisplacement()) * rec;
		}
	}
	MuSource ms = param.muSource;
	if( ms==MU_EACH_ITER || ms==MU_FIRST_ITER && beforeIter ) {
		float r = param.muProfileRatio;
		if( r == 0 ) {
			mu = 1e10;
		} else if( Fd == 0 ) {
			mu = 0;
		} else {
			mu = Fc*(1-r)/(Fd*r);
		}
	}
	std::clog 
		<< "mu=" << mu 
		<< "\t\tFc=" << Fc 
		<< "\tFd=" << Fd 
		<< std::endl;
	return Fc + mu*Fd;
}

///full task without affine (or more general) models
static void unconstrainedTask(Segmentation & s, const MotionParams & param)
{
	float mu = param.muValue;
	double F0 = targetFunctional(s, mu, true, param);
	clog << "F(initial approximation) = " << F0 << endl;
	for( int i = 0; ; ++i ) {
		if( i == param.maxIterations ) {
			clog << "maximum iterations number exceeded." << endl;
			break;
		}
		if( F0==0 ) break;
		makeIterStep( s, mu );
		float mu0 = mu;
		double F = targetFunctional(s, mu, false, param);
		clog << "F(step " << i << ") = " << F << endl;
		if( (mu<=0 && mu0<=0) || 
			(mu-mu0)/std::max(fabs(mu),fabs(mu0)) < param.minRelativeFimpr &&
			(F-F0)/fabs(F0) < param.minRelativeFimpr ) 
		{
			clog << "functional and mu values were changed too little on previous iteration." << endl;
			break;
		}
		F0 = F;
	}
}

inline static void logMotionQuality(const Segmentation &s)
{
	clog << "motion quality=" << s.motionQuality() << endl;
}

template <class T>
void Segmentation::coherentMotion(T & transform, const MotionParams & param)
{
	clog << "\nFunctional maximization" << endl;

	if (param.maxIterations >= 0) {
		intitialApproximation(*this, param.tau);
		unconstrainedTask(*this, param);
	} else {
		for (int i = 0; i != count(); ++i) {
			Segment &sg = segmentByIndex(i);
			sg.setPdisplacement(int_point(0,0));
		}
	}

	logMotionQuality(*this);

	if (count() < 10) {
		std::clog << "Segments count is less than 10, so it was decided not to find motion approximation\n";
		transform.identity();
	} else {
		int i;
		positions();

		for( i = 0; i != param.correctionLoops; ++i ) {
			determineTransform2(transform, param.transformClass);
			displaceByTransform(transform);
			logMotionQuality(*this);
			unconstrainedTask(*this, param);
			logMotionQuality(*this);
		}

		determineTransform(transform, param.transformClass);
		if (param.finalCorrection) {
			displaceByTransform(transform);
			logMotionQuality(*this);
		}
	}
}

//template void Segmentation::coherentMotion(img::QuadraticForm &, const MotionParams &);
template void Segmentation::coherentMotion(affine &, const MotionParams &);
